"""

@Time: 2021/5/8 15:22
@Author:
@File: main.py
"""
import pandas as pd
from typing import Tuple, List
from data2neo4j import DataToNeo4j

# 读取 Excel 数据到 DF 中
invoice_data: pd.DataFrame = pd.read_excel('./Invoice_data_Demo.xls', header=0)


def extract_relations(invoice_data: pd.DataFrame) -> pd.DataFrame:
    """
    将 invoice_data 中的金额、销售方名称、购买方名称三列提取出来形成一个新的 DataFrame
    """
    sell_list = []  # 所有`金额`的数据
    money_list = [] # 所有`销售方名称`数据
    buy_list = []   # 所有`购买方名称`数据
    # 将数据填入上面的 list 中
    for row in range(0, len(invoice_data)):
        money_list.append(invoice_data['金额'][row])
        sell_list.append(invoice_data['销售方名称'][row])
        buy_list.append(invoice_data['购买方名称'][row])
    # 将数据转换成 DF
    df_data = pd.DataFrame({
        'buy': buy_list,
        'money': money_list,
        'sell': sell_list
    })
    return df_data


def extract_nodes(invoice_data: pd.DataFrame) -> Tuple[List, List]:
    """
    提取节点数据
    :param invoice_data:
    :return:
    """
    # 取出购买方名称到list
    node_buy_key = []
    for i in range(0, len(invoice_data)):
        node_buy_key.append(invoice_data['购买方名称'][i])
    # 取出销售方名称到list
    node_sell_key = []
    for i in range(0, len(invoice_data)):
        node_sell_key.append(invoice_data['销售方名称'][i])
    # 去除重复的发票名称
    node_buy_key = list(set(node_buy_key))
    node_sell_key = list(set(node_sell_key))
    return node_buy_key, node_sell_key

neo4j_conn = DataToNeo4j()

relations = extract_relations(invoice_data)
node_buy_key, node_sell_key = extract_nodes(invoice_data)

neo4j_conn.create_node(node_buy_key, node_sell_key)
neo4j_conn.create_relation(relations)